

'''
直接将检测结果数据流返回，也不保存原始输入图像
'''

import base64

from PIL import Image
from io import BytesIO
from detector import PestDetector
from fastapi import FastAPI, File, UploadFile
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI
from fastapi import Request, Depends
app = FastAPI()

# 配置允许域名
origins = [ "*" ]
# origins = ["*"]  # 允许所有的请求

#  配置允许域名列表、允许方法、请求头、cookie等
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'

# 定义一个全局变量
@app.post('/predict')
async def predict(request: Request):
    '''
        {
          "data": [
            "<base64-encoded-image-data-1>",
            "<base64-encoded-image-data-2>",
            "<base64-encoded-image-data-3>",
            ...
          ]
        }
        返回的图片是base64编码
    '''
    # 从请求中检索出Base64编码的图像数据
    json_data = await request.json()
    data = json_data['data']
    # 初始化一个空数组以存储每个图像的预测结果
    predictions = []
    for item in data:
        img_name = item['name']
        image_data = item['img']
        # 解码Base64编码的图像数据并转换为NumPy数组
        decoded_data = base64.b64decode(image_data)
        image = Image.open(BytesIO(decoded_data))
        if image.mode == 'RGBA':
            image = image.convert('RGB')
        # detector = PestDetector()
        # image_pred, predicted = detector.predict(image, img_name)
        # image_pred, predicted =  app.config['detector'].predict(image, img_name)  # 使用全局变量来做
        image_pred, predicted = PestDetector().predict(image, img_name)
        # 将预测结果转换为JPEG编码的字节流
        buffer = BytesIO()
        image_pred.save(buffer, format="JPEG")
        image_bytes = buffer.getvalue()
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        # print(image_base64)
        # 将当前图像的预测结果（包括预测的图像本身）添加到预测结果数组中
        predictions.append({'imgname':img_name,'pred_info': predicted, 'image': image_base64})
    return predictions


@app.post("/image/predict")
async def upload_image(file: UploadFile = File(...)):
    # 读取上传的文件内容
    image_data = await file.read()
    image = Image.open(BytesIO(image_data))
    image_pred, predicted = PestDetector().predict(image, file.filename )
    buffer = BytesIO()
    image_pred.save(buffer, format="JPEG")
    image_bytes = buffer.getvalue()
    image_base64 = base64.b64encode(image_bytes).decode('utf-8')
    return {'imgname': file.filename , 'pred_info': predicted, 'image': image_base64}

if __name__ == '__main__':
    uvicorn.run(app =app, host = '0.0.0.0', port=8000)

